import os
import ssl
import uuid
import aiohttp
import argparse
from datetime import datetime
from fake_useragent import UserAgent
from colorama import *

green = Fore.LIGHTGREEN_EX
red = Fore.LIGHTRED_EX
magenta = Fore.LIGHTMAGENTA_EX
white = Fore.LIGHTWHITE_EX
black = Fore.LIGHTBLACK_EX
reset = Style.RESET_ALL
yellow = Fore.LIGHTYELLOW_EX

ready_proxies = []
running_proxies = []
checked_proxies = []
proxy_ips = {}


async def ipinfo(proxy):
    async with aiohttp.ClientSession() as client:
        try:
            result = await client.get("https://api.bigdatacloud.net/data/client-ip", proxy=proxy)
            return await result.json()
        except Exception as e:
            print(f"{red}get_ip_error: {white}{e}")
            return {}
class Log:
    def __init__(self, channel):
        self.channel = channel

    def info(self, msg):
        now = datetime.now().isoformat(" ").split(".")[0]
        print(f"{black}[{now}][{self.channel}]{white} {msg}{reset}")

    def log_start(self, msg):
        now = datetime.now().isoformat(" ").split(".")[0]
        print(f"{black}[{now}][{self.channel}]{white} {msg}{reset}", end="", flush=True)

    def log_append(self, msg):
        print(msg, end="", flush=True)

    def log_end(self, msg=""):
        print(msg, flush=True)

class Grass:
    def __init__(self, userid, proxy):
        self.userid = userid
        self.proxy = proxy
        self.ip = proxy_ips.get(self.proxy)
        self.logger = Log(self.ip)
        self.ses = aiohttp.ClientSession()

    def get_ssl_context(self):
        ssl_context = ssl.create_default_context()
        ssl_context.check_hostname = False
        ssl_context.verify_mode = ssl.CERT_NONE
        return ssl_context

    async def start(self):
        browser_id = uuid.uuid5(uuid.NAMESPACE_URL, self.proxy)
        useragent = UserAgent().random
        headers = {
            "Host": "proxy2.wynd.network:4650",
            "Connection": "Upgrade",
            "Pragma": "no-cache",
            "Cache-Control": "no-cache",
            "User-Agent": useragent,
            "Upgrade": "websocket",
            "Origin": "chrome-extension://lkbnfiajjmbhnfledhphioinpickokdi",
            "Sec-WebSocket-Version": "13",
            "Accept-Language": "en-US,en;q=0.9,id;q=0.8",
        }
        auth_fail = 0
        send_fail_times = 0
        send_success_times = 0
        while True:
            try:
                async with self.ses.ws_connect(
                        "wss://proxy2.wynd.network:4650/",
                        headers=headers,
                        proxy=self.proxy,
                        timeout=1000,
                        autoclose=False,
                        ssl=self.get_ssl_context()
                ) as wss:
                    res = await wss.receive_json()
                    auth_id = res.get("id")
                    if auth_id is None:
                        raise Exception('auth_id is none')
                    auth_data = {
                        "id": auth_id,
                        "origin_action": "AUTH",
                        "result": {
                            "browser_id": browser_id.__str__(),
                            "user_id": self.userid,
                            "user_agent": useragent,
                            "timestamp": int(datetime.now().timestamp()),
                            "device_type": "extension",
                            "version": "4.26.2",
                            "extension_id": "lkbnfiajjmbhnfledhphioinpickokdi",
                        },
                    }
                    await wss.send_json(auth_data)
                    self.logger.info(f"{white}{self.proxy}{green} connected{white}...")
                    if self.proxy not in running_proxies:
                        running_proxies.append(self.proxy)
                    auth_fail = 0
                    while True:
                        if wss.closed:
                            self.logger.info('socket_closed, re_auth')
                            break
                        try:
                            ping_data = {
                                "id": uuid.uuid4().__str__(),
                                "version": "1.0.0",
                                "action": "PING",
                                "data": {},
                            }
                            await wss.send_json(ping_data)
                            self.logger.log_start(f"{white}send_data...")
                            pong_data = {"id": "F3X", "origin_action": "PONG"}
                            await wss.send_json(pong_data)
                            send_success_times += 1
                            self.logger.log_end(f"{white}done[{send_success_times}]")
                            await countdown(120)
                        except Exception as e:
                            send_fail_times += 1
                            self.logger.info(f"{red}send_data_error[{send_fail_times}] : {white}{e}")
                            break
                if send_fail_times >= 3:
                    self.logger.info(f"{red}send_fail_times >= 3, skip the proxy !")
                    await self.ses.close()
                    if self.proxy in running_proxies:
                        running_proxies.remove(self.proxy)
                    return

            except KeyboardInterrupt:
                await self.ses.close()
                break
            except Exception as e:
                self.logger.info(f"{red}auth_error : {white}{e}")
                if auth_fail < 1:
                    auth_fail += 1
                    await countdown(10)
                    continue
                else:
                    self.logger.info(f"{yellow}max retrying reached, skip the proxy !")
                    await self.ses.close()
                    if self.proxy in running_proxies:
                        running_proxies.remove(self.proxy)
                    break


async def countdown(t):
    for i in range(t, 0, -1):
        minute, seconds = divmod(i, 60)
        hour, minute = divmod(minute, 60)
        seconds = str(seconds).zfill(2)
        minute = str(minute).zfill(2)
        hour = str(hour).zfill(2)
        await asyncio.sleep(1)


async def task_monitor(user_id, proxy_file_name, max_running_count):
    global checked_proxies
    """监控任务数量，如果运行中的任务数量少于最大值，则动态添加任务"""
    logger = Log('task_monitor')
    logger.info(f'task_monitor started')
    while True:
        logger.info(f'task_monitor_check - {len(running_proxies)} / {max_running_count}{reset}')
        if len(running_proxies) >= max_running_count:
            await countdown(60)
            continue
        proxies = open(proxy_file_name, "r").read().splitlines()
        if len(proxies) == 0:
            logger.info(f"{red} {proxy_file_name} is empty")
        if len(ready_proxies) == 0:
            # 检查代理
            no_valid = True
            for proxy in proxies:
                if proxy in running_proxies:
                    continue
                if proxy in checked_proxies:
                    continue
                no_valid = False
                ip = (await ipinfo(proxy)).get('ipString')
                checked_proxies.append(proxy)
                if ip is not None:
                    ready_proxies.append(proxy)
                    proxy_ips[proxy] = ip
                    logger.info(f"{white}proxy[{proxy}] {green}is valid [{ip}]{reset}")
                    break
                else:
                    logger.info(f"{white}proxy[{proxy}] {red}is invalid{reset}")
            if no_valid:
                logger.info(f'{red}no valid proxy, reset used_proxy')
                checked_proxies = []
                await countdown(2)
                continue

        if len(ready_proxies) == 0:
            await countdown(60)
            continue

        proxy = ready_proxies.pop(0)  # 从待执行任务队列中取出任务
        logger.info(f"{green}started a new task on proxy [{proxy}]")
        asyncio.create_task(Grass(user_id, proxy).start())  # 启动任务
        await countdown(5)


async def main():
    arg = argparse.ArgumentParser()
    arg.add_argument(
        "--user_id", "-U", help="user_id generated by setup.py", required=True
    )
    args = arg.parse_args()
    userid = args.user_id
    if userid is None or len(userid) <= 0:
        print(f"{red}错误 : {white}请先输入您的用户ID!")
        exit()
    proxy_path = f"proxies_{userid}.txt"
    if not os.path.exists(proxy_path):
        print(f"{red}{proxy_path} 未找到，请确保 {proxy_path} 可用！")
        exit()

    # 启动任务监控
    await asyncio.create_task(task_monitor(userid, proxy_path, 30))


if __name__ == "__main__":
    try:
        import asyncio

        if os.name == "nt":
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
        asyncio.run(main())
    except KeyboardInterrupt:
        exit()
